/*******************************************************************/
/*                                                                 */
/*  NAME    = send_pcap.cpp                                                   */
/*  NOTE    =                                                      */
/*  DATE    = 2014/12/10 by yinghl                                    */
/*******************************************************************/

#include "send_pcap.h"

#include <iostream>
#include <sstream>
#include <algorithm>

using namespace std;

/*******************************************************************************
*  FUNC     :   构造函数初始化一些成员变量
*  ARGS     :   pcap  (pcap文件)
                fps (发送包的频率)
                libnet (libnet句柄)
                changePortCmpData (更改端口号的条件)
                changePort (更改的端口号)
*  RTN      :
*  NOTE     :
*******************************************************************************/
SendPcap::SendPcap(const string &pcap, uint32_t fps, libnet_t *libnet,
        const string &changePortCmpData, int changePort)
    : pcap_(pcap), pcapHandle_(NULL), libnetHandle_(libnet),
    fps_(fps), changePortCmpData_(changePortCmpData), changePort_(changePort)
{
}

/********************************************************************************
*  FUNC     :   析构函数释放pcap/libnet句柄
*  ARGS     :
*  RTN      :
*  NOTE     :
*******************************************************************************/
SendPcap::~SendPcap()
{
    if(pcapHandle_)
        pcap_close(pcapHandle_);
}

/********************************************************************************
*  FUNC     :   线程处理函数
*  ARGS     :
*  RTN      :
*  NOTE     :   解析pcap文件并发送从pcap文件中解析出来的数据包
*******************************************************************************/
void SendPcap::doSend()
{
    char errBuf[PCAP_ERRBUF_SIZE];

    //打开pcap文件
    pcapHandle_ = pcap_open_offline(pcap_.c_str(), errBuf);
    if(!pcapHandle_)
    {
        cout << "pcap_open_offline failed, " << errBuf << endl;
        return;
    }

    int status;
    const u_char *pktData = NULL;
    pcap_pkthdr *pktHeader = NULL;
    timeval lastTm = {0, 0};
    timeval st_time_start,st_time_end;
    int32_t i_send_time=0;

    static int idx = 1;
    long l_conts1000_per = 0;
    long i_time_per = 1000*1000 /1000;
    int i_flag_wait=0;
    
    if(fps_ != 0)
        l_conts1000_per = fps_ / 1000;
    gettimeofday (&st_time_start, NULL);

    do
    {
        i_send_time = 0;
        //if(fps_ && (idx % l_conts1000_per == 0) )
        //{
        //    gettimeofday (&st_time_start, NULL);
        //    i_flag_wait = 1;
        //}
        //else
            gettimeofday (&st_time_start, NULL);

        status = pcap_next_ex(pcapHandle_, &pktHeader, &pktData);
        if(status == 1 && pktHeader && pktData)
        {
            //不是压力测试
            if(fps_ == 0)
            {
                if(lastTm.tv_sec == 0 && lastTm.tv_usec == 0)
                    lastTm = pktHeader->ts;
                else
                {
                    int32_t pktTimeLev = 0;
                    if(lastTm.tv_sec != 0 || lastTm.tv_usec != 0)
                    {
                        pktTimeLev = (pktHeader->ts.tv_sec - lastTm.tv_sec) * 1000000
                            + (pktHeader->ts.tv_usec - lastTm.tv_usec);
                        //cout << "pktTimeLev: " << pktTimeLev << endl;
                        if(pktTimeLev < 0) pktTimeLev = 0;
                    }
    
                    lastTm = pktHeader->ts;
                    
                    gettimeofday (&st_time_end, NULL);
                    i_send_time = (st_time_end.tv_sec - st_time_start.tv_sec) * 1000000
                                + (st_time_end.tv_usec - st_time_start.tv_usec);
                    // 2015/03/23 发包时间不控制
                    //if(pktTimeLev - i_send_time > 0)
                    //{
                    //    printf("idx1 =%d usleep %d\n", idx-1, pktTimeLev - i_send_time);
                    //    usleep(pktTimeLev - i_send_time);
                    //}
                }
            }

            //解析pcap包数据
            parsePcapData(pktHeader->caplen, pktData);
            //cout << "idx: " << idx << " done" << endl;
        }

        if(status == 1)
        {
            //cout << "sleepTime: " << l_sleepTime << "us" << endl;
            //压力测试时，控制发包速度
            if(fps_ && i_flag_wait)
            {
                gettimeofday (&st_time_end, NULL);
                i_send_time = (st_time_end.tv_sec - st_time_start.tv_sec) * 1000000
                            + (st_time_end.tv_usec - st_time_start.tv_usec);
                
                // 2015/03/23 发包时间不控制
                //if(i_time_per - i_send_time > 0)
                //{
                //    printf("idx2 =%d usleep %ld\n", idx-1, i_time_per - i_send_time);
                //    usleep(i_time_per - i_send_time);
                //}

                i_flag_wait = 0;
            }
                
            idx++;
        }
    } while(status == 1);

    //cout << "send " << idx-1 << " packages" << endl;
}

/********************************************************************************
*  FUNC     :   解析pcap包数据
*  ARGS     :   capLen (pcap数据包长度)
                pktData (pcap数据包数据)
*  RTN      :
*  NOTE     :
*******************************************************************************/
void SendPcap::parsePcapData(u_int capLen, const u_char *pktData)
{
    ether_header ethHeader;
    u_int linkOffset;
    if(!parseEtherHeader(capLen, pktData, ethHeader, linkOffset))
    {
        cout << "parseEtherHeader failed" << endl;
        return;
    }
    //cout << "linkOffset: " << linkOffset << endl;
    iphdr ipHeader;
    const u_char *ipHeaderData = NULL;
    if(!parseIPHeader(capLen, pktData, linkOffset, ipHeader, &ipHeaderData))
    {
        cout << "parseIPHeader failed" << endl;
        return;
    }

    if(ipHeader.protocol == IPPROTO_TCP)
    {
        parseTCPProtol(capLen, pktData, linkOffset, ethHeader, ipHeader);
    }
    else if(ipHeader.protocol == IPPROTO_IP)
    {
        parseIPProtol(capLen, pktData, linkOffset, ethHeader, ipHeader);
    }
    else
    {
        cout << "ip protocol not support, " << int(ipHeader.protocol) << endl;
    }
}

/********************************************************************************
*  FUNC     :   解析mac头
*  ARGS     :   capLen (pcap数据包长度)
                pktData (pcap数据包数据)
                ethHeader (成功返回后保存ether信息)
                linkOffset (成功返回后保存linkOffset)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::parseEtherHeader(u_int capLen, const u_char *pktData,
        ether_header &ethHeader, u_int &linkOffset)
{
    if(!pktData) return false;

    uint ethHeaderLen = sizeof(ether_header);
    if(capLen < ethHeaderLen)  //小于最小长度
        return false;

    //获取mac层长度
    if(pktData[12] == 0x8 && pktData[13] == 0x0) /* Regular ethernet */
        linkOffset = 14;
    else if(pktData[12] == 0x81 && pktData[13] == 0x0) /* Skip 802.1Q VLAN and priority information */
        linkOffset = 18;
    else    /* non-ip frame */
        return false;

    memset(&ethHeader, 0, ethHeaderLen);
    memcpy(&ethHeader, pktData, ethHeaderLen);

    if(ntohs(ethHeader.ether_type) != ETHERTYPE_IP)
        return false;

    return true;
}

/********************************************************************************
*  FUNC     :   解析ip头
*  ARGS     :   capLen (pcap数据包长度)
                pktData (pcap数据包数据)
                linkOffset (mac层长度)
                ipHeader (成功返回后保存ip头信息)
                ipHeaderData (成功返回后保存ip头位置)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::parseIPHeader(u_int capLen, const u_char *pktData,
        u_int linkOffset, iphdr &ipHeader, const u_char **ipHeaderData)
{
    if(!pktData) return false;

    if(capLen < linkOffset + sizeof(iphdr)) //无法获取ip头长度不够
        return false;

    *ipHeaderData = pktData + linkOffset; //ip数据起始位置
    memset(&ipHeader, 0, sizeof(iphdr));
    memcpy(&ipHeader, *ipHeaderData, sizeof(iphdr));

    return true;
}

/********************************************************************************
*  FUNC     :   解析tcp头
*  ARGS     :   capLen (pcap数据包长度)
                pktData (pcap数据包数据)
                linkOffset (mac层长度)
                ipHeader (ip头)
                tcpHeaderData (成功返回后保存tcp头位置)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::parseTCPHeader(u_int capLen, const u_char *pktData,
        u_int linkOffset, const iphdr &ipHeader, tcphdr &tcpHeader,
        const u_char **tcpHeaderData)
{
    if(!pktData) return false;

    if(capLen < linkOffset + 4*ipHeader.ihl + sizeof(tcphdr)) //无法获取tcp头长度不够
        return false;

    *tcpHeaderData = pktData + linkOffset + 4*ipHeader.ihl;
    memset(&tcpHeader, 0, sizeof(tcphdr));
    memcpy(&tcpHeader, *tcpHeaderData, sizeof(tcphdr));

    return true;
}

/********************************************************************************
*  FUNC     :   解析tcp数据长度
*  ARGS     :   capLen (pcap数据包长度)
                linkOffset (mac层长度)
                ipHeader (ip头信息)
*  RTN      :   数据长度
*  NOTE     :
*******************************************************************************/
int SendPcap::parseTCPDataLen(u_int capLen, u_int linkOffset,
        const iphdr &ipHeader, const tcphdr &tcpHeader)
{

    int tcpDataLen = 0;
    if(ntohs(ipHeader.tot_len) == 0)
    {
        tcpDataLen = capLen - linkOffset - 4 * ipHeader.ihl - 4 * tcpHeader.doff;
    }
    else
    {
        tcpDataLen = ntohs(ipHeader.tot_len) - 4 * ipHeader.ihl - 4 * tcpHeader.doff;
    }

    return tcpDataLen;
}

/********************************************************************************
*  FUNC     :   解析tcp协议数据
*  ARGS     :   capLen (pcap数据包长度)
                pktData (pcap数据)
                linkOffset (mac层长度)
                ethHeader (网卡信息)
                ipHeader (ip头信息)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::parseTCPProtol(u_int capLen, const u_char *pktData, u_int linkOffset,
        const ether_header &ethHeader, const iphdr &ipHeader)
{
    if(!pktData) return false;

    tcphdr tcpHeader;
    const u_char *tcpHeaderData = NULL;
    if(!parseTCPHeader(capLen, pktData, linkOffset, ipHeader, tcpHeader, &tcpHeaderData))
    {
        cout << "parseTCPHeader failed" << endl;
        return false;
    }

    //saddr
    if(needChangePort("tcp", ipHeader.saddr, ntohs(tcpHeader.source)))
    {
        tcpHeader.source = htons(changePort_);
    }

    //daddr
    if(needChangePort("tcp", ipHeader.daddr, ntohs(tcpHeader.dest)))
    {
        tcpHeader.dest = htons(changePort_);
    }

    int tcpDataLen = parseTCPDataLen(capLen, linkOffset, ipHeader, tcpHeader);
    //cout << "tcpDataLen: " << tcpDataLen << endl;
    if(tcpDataLen < 0)
    {
        cout << "warn:please check package" << endl;
        return false;
    }
    //cout << "tcpDataLen: " << tcpDataLen << endl;
    const u_char *tcpData = tcpHeaderData + 4 * tcpHeader.doff;

    int tcpOptionLen = 4 * tcpHeader.doff - 20;
    const u_char *tcpOptionData = NULL;
    if(tcpOptionLen > 0)
    {
        tcpOptionData = tcpHeaderData + sizeof(tcphdr);
    }
    else
    {
        tcpOptionLen = 0;
    }
    //cout << "tcpOptionLen: " << tcpOptionLen << endl;

    return sendTCPProtol(ethHeader, ipHeader, tcpHeader, tcpOptionData, tcpOptionLen,
            tcpData, tcpDataLen);
}


/********************************************************************************
*  FUNC     :   发送tcp协议数据包
*  ARGS     :   ethHeader (网卡信息)
                ipHeader (ip头信息)
                tcpHeader (tcp头信息)
                tcpOptionData (tcp option数据)
                tcpOptionLen (tcp option数据长度)
                tcpData (tcp数据)
                tcpDataLen (tcp数据长度)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::sendTCPProtol(const ether_header &ethHeader, const iphdr &ipHeader,
        const tcphdr &tcpHeader, const u_char *tcpOptionData, int tcpOptionLen,
        const u_char *tcpData, int tcpDataLen)
{
    if(!initLibNet()) return false;

    //cout << "libnet_init/libnet_clear_packet" << endl;
    if(!buildTCPOptions(tcpOptionData, tcpOptionLen))
        return false;
    //cout << "buildTCPOptions tcpOptionLen: " << tcpOptionLen << endl;

    if(!buildTCPData(tcpHeader, tcpData, tcpDataLen, tcpOptionLen))
        return false;
    //cout << "buildTCPData tcpDataLen: " << tcpDataLen << endl;

    if(!buildIPv4(ipHeader, NULL, 0))
        return false;
    //cout << "buildIPv4" << endl;

    if(!buildEther(ethHeader))
        return false;
    //cout << "buildEther" << endl;

    return libNetWrite();
}

/********************************************************************************
*  FUNC     :   tcp options数据组包
*  ARGS     :   tcpOptionData (tcp option数据)
                tcpOptionLen (tcp option数据长度)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::buildTCPOptions(const u_char *tcpOptionData, int tcpOptionLen)
{
    if(tcpOptionLen > 0 && tcpOptionData)
    {
        libnet_ptag_t ret = libnet_build_tcp_options(tcpOptionData, tcpOptionLen,
                libnetHandle_, 0);
        if(ret == -1)
        {
            cout << "libnet_build_tcp_options failed, " << libnet_geterror(libnetHandle_) << endl;
            return false;
        }
    }

    return true;
}

/********************************************************************************
*  FUNC     :   tcp数据组包
*  ARGS     :   tcpHeader (tcp头信息)
                tcpData (tcp数据)
                tcpDataLen (tcp数据长度)
                tcpOptionLen (tcp options长度)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::buildTCPData(const tcphdr &tcpHeader, const u_char *tcpData, int tcpDataLen,
        int tcpOptionLen)
{
    libnet_ptag_t ret = libnet_build_tcp(
            ntohs(tcpHeader.source),                              /* source port */
            ntohs(tcpHeader.dest),                                /* destination port */
            ntohl(tcpHeader.seq ),                                /* sequence number */
            ntohl(tcpHeader.ack_seq ),                            /* acknowledgement num */
            (tcpHeader.urg<<5)|(tcpHeader.ack<<4)|(tcpHeader.psh<<3)|(tcpHeader.rst<<2)|(tcpHeader.syn<<1)|(tcpHeader.fin<<0),                                            /* control flags */
            ntohs(tcpHeader.window ),                             /* window size */
            ntohs(tcpHeader.check ),                              /* checksum */
            ntohs(tcpHeader.urg_ptr ),                            /* urgent pointer */
            LIBNET_TCP_H + tcpOptionLen + tcpDataLen,             /* TCP packet size */
            (uint8_t*)tcpData,                                     /* payload */
            tcpDataLen,                                          /* payload size */
            libnetHandle_,                                          /* libnet handle */
            0);
    if(ret == -1)
    {
        cout << "libnet_build_tcp failed, " << libnet_geterror(libnetHandle_) << endl;
        return false;
    }

    return true;
}

/********************************************************************************
*  FUNC     :   IPv4数据组包
*  ARGS     :   ipHeader (ip头信息)
*               ipData (ip数据)
*               ipDataLen (ip数据长度)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::buildIPv4(const iphdr &ipHeader, const u_char *ipData, int ipDataLen)
{
    libnet_ptag_t ret = libnet_build_ipv4(
            //LIBNET_IPV4_H + LIBNET_TCP_H + ui_payload_len,
            ntohs(ipHeader.tot_len),                      /* length */
            ipHeader.tos,                                 /* TOS */
            ntohs(ipHeader.id),                           /* IP ID */
            ntohs(ipHeader.frag_off),                     /* IP Frag */
            ipHeader.ttl,                                 /* TTL */
            ipHeader.protocol,  //IPPROTO_TCP,            /* protocol */
            ntohs(ipHeader.check),                        /* checksum */
            ipHeader.saddr,                               /* source IP */
            ipHeader.daddr,                               /* destination IP */
            ipData,                                            /* payload */
            ipDataLen,                                               /* payload size */
            libnetHandle_,                                  /* libnet handle */
            0);                                              /* libnet id */

    if(ret == -1)
    {
        cout << "libnet_build_ipv4 failed, " << libnet_geterror(libnetHandle_) << endl;
        return false;
    }

    return true;
}

/********************************************************************************
*  FUNC     :   mac数据组包
*  ARGS     :   ethHeader (mac头信息)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::buildEther(const ether_header &ethHeader)
{
    libnet_ptag_t ret = libnet_build_ethernet(
        ethHeader.ether_dhost,                       /* ethernet dest */
        ethHeader.ether_shost,                       /* ethernet source */
        ETHERTYPE_IP,                               /* protocol type */
        NULL,                                           /* payload */
        0,                                              /* payload size */
        libnetHandle_,                                 /* libnet handle */
        0);                                             /* libnet id */
    if(ret == -1)
    {
        cout << "libnet_build_ethernet failed, " << libnet_geterror(libnetHandle_) << endl;
        return false;
    }

    return true;
}

/********************************************************************************
*  FUNC     :   解析ip协议数据
*  ARGS     :   capLen (pcap数据包长度)
                pktData (pcap数据)
                linkOffset (mac层长度)
                ethHeader (网卡信息)
                ipHeader (ip头信息)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::parseIPProtol(u_int capLen, const u_char *pktData,
        u_int linkOffset, const ether_header &ethHeader,
        const iphdr &ipHeader)
{
    if(!pktData) return false;

    int ipDataLen = parseIPDataLen(capLen, linkOffset, ipHeader);
    if(ipDataLen < 0)
    {
        cout << "warn:please check package" << endl;
        return false;
    }

    //这里可能有options
    //int optionsLen = 4 * ipHeader.ihl - LIBNET_IPV4_H;

    const u_char *ipData = pktData + linkOffset + 4 * ipHeader.ihl;

    return sendIPProtol(ethHeader, ipHeader, ipData, ipDataLen);
}

/********************************************************************************
*  FUNC     :   解析ip数据长度
*  ARGS     :   capLen (pcap数据包长度)
                linkOffset (mac层长度)
                ipHeader (ip头信息)
*  RTN      :   数据长度
*  NOTE     :
*******************************************************************************/
int SendPcap::parseIPDataLen(u_int capLen, u_int linkOffset,
        const iphdr &ipHeader)
{
    int ipDataLen = 0;
    if(ntohs(ipHeader.tot_len) == 0)
    {
        ipDataLen = capLen - linkOffset - 4 * ipHeader.ihl;
    }
    else
    {
        ipDataLen = ntohs(ipHeader.tot_len) - 4 * ipHeader.ihl;
    }

    return ipDataLen;
}

/********************************************************************************
*  FUNC     :   发送ip协议数据包
*  ARGS     :   ethHeader (网卡信息)
                ipHeader (ip头信息)
                ipData (ip数据)
                ipDataLen (ip数据长度)
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::sendIPProtol(const ether_header &ethHeader, const iphdr &ipHeader,
        const u_char *ipData, int ipDataLen)
{
    if(!initLibNet()) return false;

    if(!buildIPv4(ipHeader, ipData, ipDataLen))
        return false;

    if(!buildEther(ethHeader))
        return false;

    return libNetWrite();
}

/********************************************************************************
*  FUNC     :   初始化libnet
*  ARGS     :
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::initLibNet()
{
    if(!libnetHandle_)
    {
        cout << "libnetHandle_ is NULL" << endl;
        return false;
    }
    else
    {
        libnet_clear_packet(libnetHandle_);
    }

    return true;
}

/********************************************************************************
*  FUNC     :   libnet发送数据
*  ARGS     :
*  RTN      :   成功: true, 失败: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::libNetWrite()
{
    libnet_ptag_t ret = libnet_write(libnetHandle_);
    if(ret != -1)
    {
        //cout << "libnet_write " << ret << " bytes"<< endl;
    }
    else
    {
        cout << "libnet_write failed, " << libnet_geterror(libnetHandle_) << endl;
    }

    return (ret != -1);
}

/********************************************************************************
*  FUNC     :   检查是否需要转换端口
*  ARGS     :   flag (协议：tcp/udp)
                addr (ip地址)
                port (端口号)
*  RTN      :   需要转换: true, 不需要转换: false
*  NOTE     :
*******************************************************************************/
bool SendPcap::needChangePort(const std::string &flag, uint32_t addr, int port)
{
    const u_char *strAddr = (const u_char*)&addr;
    ostringstream oss;
    oss << flag << ":" << int(strAddr[0]) << "." << int(strAddr[1])
        << "." << int(strAddr[2]) << "." << int(strAddr[3]) << ":" << port;
    string tmpCmpData = oss.str();

    transform(tmpCmpData.begin(), tmpCmpData.end(), tmpCmpData.begin(), ::toupper);
    transform(changePortCmpData_.begin(), changePortCmpData_.end(), changePortCmpData_.begin(),
            ::toupper);

    return tmpCmpData == changePortCmpData_;
}
